#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
压力测试运行脚本

提供便捷的命令行接口来执行不同场景的压力测试
支持多种测试场景和环境配置

使用示例:
    python tests/performance/run_performance_test.py --scenario load --environment local
    python tests/performance/run_performance_test.py --scenario stress --users 200 --time 300s
    python tests/performance/run_performance_test.py --help
"""

import argparse
import os
import sys
import subprocess
import json
import time
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from tests.performance.config import (
    TEST_SCENARIOS, 
    ENVIRONMENT_CONFIGS, 
    get_config,
    REPORT_CONFIG
)
from utils.logger import get_logger
from utils.helpers import ensure_dir_exists, get_timestamp

# 配置日志
logger = get_logger(__name__)


class PerformanceTestRunner:
    """
    压力测试运行器
    
    负责执行Locust压力测试并生成报告
    """
    
    def __init__(self):
        self.project_root = project_root
        self.locustfile_path = self.project_root / "tests" / "performance" / "locustfile.py"
        self.report_dir = self.project_root / REPORT_CONFIG["output_dir"]
        
        # 确保报告目录存在
        ensure_dir_exists(str(self.report_dir))
    
    def check_dependencies(self) -> bool:
        """
        检查依赖项是否安装
        
        Returns:
            bool: 依赖项是否满足
        """
        try:
            import locust
            logger.info(f"Locust版本: {locust.__version__}")
            return True
        except ImportError:
            logger.error("Locust未安装，请运行: pip install locust")
            return False
    
    def validate_config(self, scenario: str, environment: str) -> bool:
        """
        验证配置参数
        
        Args:
            scenario: 测试场景
            environment: 测试环境
        
        Returns:
            bool: 配置是否有效
        """
        if scenario not in TEST_SCENARIOS:
            logger.error(f"无效的测试场景: {scenario}")
            logger.info(f"可用场景: {', '.join(TEST_SCENARIOS.keys())}")
            return False
        
        if environment not in ENVIRONMENT_CONFIGS:
            logger.error(f"无效的测试环境: {environment}")
            logger.info(f"可用环境: {', '.join(ENVIRONMENT_CONFIGS.keys())}")
            return False
        
        return True
    
    def build_locust_command(self, 
                           scenario: str,
                           environment: str,
                           users: Optional[int] = None,
                           spawn_rate: Optional[int] = None,
                           run_time: Optional[str] = None,
                           headless: bool = True,
                           html_report: bool = True,
                           csv_report: bool = True) -> List[str]:
        """
        构建Locust命令
        
        Args:
            scenario: 测试场景
            environment: 测试环境
            users: 用户数量
            spawn_rate: 用户生成速率
            run_time: 运行时间
            headless: 是否无界面运行
            html_report: 是否生成HTML报告
            csv_report: 是否生成CSV报告
        
        Returns:
            List[str]: Locust命令参数列表
        """
        config = get_config(scenario, environment)
        
        # 使用传入参数或配置默认值
        final_users = users or config.users
        final_spawn_rate = spawn_rate or config.spawn_rate
        final_run_time = run_time or config.run_time
        
        # 生成报告文件名
        timestamp = get_timestamp()
        report_prefix = f"{scenario}_{environment}_{timestamp}"
        
        # 构建基础命令
        cmd = [
            "locust",
            "-f", str(self.locustfile_path),
            "--host", config.host,
            "--users", str(final_users),
            "--spawn-rate", str(final_spawn_rate),
            "--run-time", final_run_time
        ]
        
        # 添加无界面模式
        if headless:
            cmd.append("--headless")
        
        # 添加HTML报告
        if html_report:
            html_file = self.report_dir / f"{report_prefix}.html"
            cmd.extend(["--html", str(html_file)])
        
        # 添加CSV报告
        if csv_report:
            csv_prefix = self.report_dir / report_prefix
            cmd.extend(["--csv", str(csv_prefix)])
        
        # 添加日志级别
        cmd.extend(["--loglevel", "INFO"])
        
        return cmd
    
    def run_test(self, 
                scenario: str,
                environment: str,
                users: Optional[int] = None,
                spawn_rate: Optional[int] = None,
                run_time: Optional[str] = None,
                headless: bool = True) -> bool:
        """
        运行压力测试
        
        Args:
            scenario: 测试场景
            environment: 测试环境
            users: 用户数量
            spawn_rate: 用户生成速率
            run_time: 运行时间
            headless: 是否无界面运行
        
        Returns:
            bool: 测试是否成功
        """
        logger.info("=" * 60)
        logger.info("开始压力测试")
        logger.info(f"场景: {scenario}")
        logger.info(f"环境: {environment}")
        logger.info(f"用户数: {users or TEST_SCENARIOS[scenario].users}")
        logger.info(f"生成速率: {spawn_rate or TEST_SCENARIOS[scenario].spawn_rate}")
        logger.info(f"运行时间: {run_time or TEST_SCENARIOS[scenario].run_time}")
        logger.info("=" * 60)
        
        # 构建命令
        cmd = self.build_locust_command(
            scenario=scenario,
            environment=environment,
            users=users,
            spawn_rate=spawn_rate,
            run_time=run_time,
            headless=headless
        )
        
        logger.info(f"执行命令: {' '.join(cmd)}")
        
        try:
            # 记录开始时间
            start_time = time.time()
            
            # 执行Locust命令
            result = subprocess.run(
                cmd,
                cwd=str(self.project_root),
                capture_output=True,
                text=True,
                timeout=3600  # 1小时超时
            )
            
            # 记录结束时间
            end_time = time.time()
            duration = end_time - start_time
            
            logger.info(f"测试完成，耗时: {duration:.2f}秒")
            
            if result.returncode == 0:
                logger.info("压力测试执行成功")
                logger.info("标准输出:")
                logger.info(result.stdout)
                
                # 分析测试结果
                self._analyze_results(scenario, environment)
                
                return True
            else:
                logger.error("压力测试执行失败")
                logger.error(f"返回码: {result.returncode}")
                logger.error("标准错误:")
                logger.error(result.stderr)
                return False
                
        except subprocess.TimeoutExpired:
            logger.error("压力测试执行超时")
            return False
        except Exception as e:
            logger.error(f"压力测试执行异常: {e}")
            return False
    
    def _analyze_results(self, scenario: str, environment: str):
        """
        分析测试结果
        
        Args:
            scenario: 测试场景
            environment: 测试环境
        """
        logger.info("分析测试结果...")
        
        # 查找最新的报告文件
        report_files = list(self.report_dir.glob(f"{scenario}_{environment}_*.html"))
        if report_files:
            latest_report = max(report_files, key=lambda x: x.stat().st_mtime)
            logger.info(f"HTML报告已生成: {latest_report}")
        
        csv_files = list(self.report_dir.glob(f"{scenario}_{environment}_*_stats.csv"))
        if csv_files:
            latest_csv = max(csv_files, key=lambda x: x.stat().st_mtime)
            logger.info(f"CSV报告已生成: {latest_csv}")
            
            # 简单分析CSV数据
            try:
                import pandas as pd
                df = pd.read_csv(latest_csv)
                
                logger.info("测试结果摘要:")
                logger.info(f"总请求数: {df['Request Count'].sum()}")
                logger.info(f"失败请求数: {df['Failure Count'].sum()}")
                logger.info(f"平均响应时间: {df['Average Response Time'].mean():.2f}ms")
                logger.info(f"最大响应时间: {df['Max Response Time'].max():.2f}ms")
                
            except ImportError:
                logger.warning("pandas未安装，跳过详细分析")
            except Exception as e:
                logger.warning(f"分析CSV文件失败: {e}")
    
    def list_scenarios(self):
        """
        列出所有可用的测试场景
        """
        print("\n可用的测试场景:")
        print("=" * 40)
        for scenario, config in TEST_SCENARIOS.items():
            print(f"{scenario:12} - {config.users:3d}用户, {config.run_time:>6}, 最大响应时间{config.max_response_time:>6.0f}ms")
    
    def list_environments(self):
        """
        列出所有可用的测试环境
        """
        print("\n可用的测试环境:")
        print("=" * 40)
        for env, config in ENVIRONMENT_CONFIGS.items():
            print(f"{env:8} - {config['host']} ({config['description']})")


def main():
    """
    主函数
    """
    parser = argparse.ArgumentParser(
        description="API压力测试运行器",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  %(prog)s --scenario load --environment local
  %(prog)s --scenario stress --users 200 --time 300s
  %(prog)s --list-scenarios
  %(prog)s --list-environments
        """
    )
    
    parser.add_argument(
        "--scenario", 
        choices=list(TEST_SCENARIOS.keys()),
        default="load",
        help="测试场景 (默认: load)"
    )
    
    parser.add_argument(
        "--environment", 
        choices=list(ENVIRONMENT_CONFIGS.keys()),
        default="local",
        help="测试环境 (默认: local)"
    )
    
    parser.add_argument(
        "--users", 
        type=int,
        help="用户数量 (覆盖场景默认值)"
    )
    
    parser.add_argument(
        "--spawn-rate", 
        type=int,
        help="用户生成速率 (覆盖场景默认值)"
    )
    
    parser.add_argument(
        "--time", 
        help="运行时间，如: 60s, 5m, 1h (覆盖场景默认值)"
    )
    
    parser.add_argument(
        "--gui", 
        action="store_true",
        help="启用Web界面 (默认无界面运行)"
    )
    
    parser.add_argument(
        "--list-scenarios", 
        action="store_true",
        help="列出所有可用的测试场景"
    )
    
    parser.add_argument(
        "--list-environments", 
        action="store_true",
        help="列出所有可用的测试环境"
    )
    
    args = parser.parse_args()
    
    # 创建测试运行器
    runner = PerformanceTestRunner()
    
    # 处理列表命令
    if args.list_scenarios:
        runner.list_scenarios()
        return
    
    if args.list_environments:
        runner.list_environments()
        return
    
    # 检查依赖项
    if not runner.check_dependencies():
        sys.exit(1)
    
    # 验证配置
    if not runner.validate_config(args.scenario, args.environment):
        sys.exit(1)
    
    # 运行测试
    success = runner.run_test(
        scenario=args.scenario,
        environment=args.environment,
        users=args.users,
        spawn_rate=args.spawn_rate,
        run_time=args.time,
        headless=not args.gui
    )
    
    if success:
        logger.info("压力测试完成")
        sys.exit(0)
    else:
        logger.error("压力测试失败")
        sys.exit(1)


if __name__ == "__main__":
    main()